/*
 * @file arma_vecmap.h
 *
 * The Map from vector to values
 *
 * @author Dahua Lin
 */

#ifndef _MSC_VER
#pragma once
#endif

#ifndef GMLPP_ARMA_VECMAP_H_
#define GMLPP_ARMA_VECMAP_H_

#include "arma_base.h"

namespace gml
{
	/**
	 * column wise vector pair evaluation
	 */
	template<class VecFun, typename T1, typename TR>
	void colwise_vecmap(VecFun fun, const Mat<T1>& X1, Mat<TR>& r)
	{
		uword n = X1.n_cols;

		TR *rp = r.memptr();
		for (uword j = 0; j < n; ++j)
		{
			rp[j] = fun(X1.unsafe_col(j));
		}
	}

	template<class VecFun, typename T1>
	inline Row<typename VecFun::result_type> colwise_vecmap(VecFun fun, const Mat<T1>& X1)
	{
		Row<typename VecFun::result_type> r(X1.n_cols);
		colwise_vecmap(fun, X1, r);
		return r;
	}


	template<class VecFun, typename T1, typename T2, typename TR>
	void colwise_vecmap2(VecFun fun, const Mat<T1>& X1, const Mat<T2>& X2, Mat<TR>& r)
	{

		check_arg( is_same_size(X1, X2),
				"dimension mismatch for colwise_vecpair_eval");
		uword n = X1.n_cols;

		TR *rp = r.memptr();
		for (uword j = 0; j < n; ++j)
		{
			rp[j] = fun(X1.unsafe_col(j), X2.unsafe_col(j));
		}
	}

	template<class VecFun, typename T1, typename T2>
	inline Row<typename VecFun::result_type> colwise_vecmap2(VecFun fun, const Mat<T1>& X1, const Mat<T2>& X2)
	{
		Row<typename VecFun::result_type> r(X1.n_cols);
		colwise_vecmap2(fun, X1, X2, r);
		return r;
	}


	/**
	 * pairwise vector pair evaluation
	 */
	template<class VecFun, typename T1, typename T2, typename TR>
	void pairwise_vecmap(VecFun fun, const Mat<T1>& X, const Mat<T2>& Y, Mat<TR>& R)
	{
		uword nx = X.n_cols;
		uword ny = Y.n_cols;
		check_arg(R.n_rows == nx && R.n_cols == ny,
				"dimension mismatch for pairwise_vecpair_eval");

		for (uword j = 0; j < ny; ++j)
		{
			for (uword i = 0; i < nx; ++i)
			{
				R.at(i, j) = fun(X.unsafe_col(i), Y.unsafe_col(j));
			}
		}
	}

	template<class VecFun, typename T1, typename T2>
	inline Mat<typename VecFun::result_type> pairwise_vecmap(VecFun fun, const Mat<T1>& X, const Mat<T2>& Y)
	{
		Mat<typename VecFun::result_type> R(X.n_cols, Y.n_cols);
		pairwise_vecmap(fun, X, Y, R);
		return R;
	}

}

#endif /* ARMA_VECMAP_H_ */
